# main.py
import asyncio
import json
from typing import Any, Dict
from pathlib import Path

import paramiko
from fastapi import FastAPI, WebSocket
from fastapi.staticfiles import StaticFiles

app = FastAPI()

# ----------- 工具函数 -----------
def test_ssh(host: str, port: int, user: str, pwd: str) -> Dict[str, Any]:
    """测试 SSH 连接，成功返回 {"ok": true}，失败返回 {"ok": false, "err": str}"""
    try:
        with paramiko.SSHClient() as ssh:
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(hostname=host, port=port, username=user, password=pwd, timeout=5)
            return {"ok": True}
    except Exception as e:
        return {"ok": False, "err": str(e)}


def list_docker_images(host: str, port: int, user: str, pwd: str):
    """列出远端镜像，返回 ["image1:latest", ...]"""
    cmd = "docker images --format '{{.Repository}}:{{.Tag}}'"
    try:
        with paramiko.SSHClient() as ssh:
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(hostname=host, port=port, username=user, password=pwd)
            stdin, stdout, stderr = ssh.exec_command(cmd)
            err = stderr.read().decode()
            if err:
                return {"ok": False, "err": err}
            images = [line.strip() for line in stdout if line.strip()]
            # 过滤掉包含<none>的镜像并排序
            images = [img for img in images if "<none>" not in img]
            images.sort()
            return {"ok": True, "images": images}
    except Exception as e:
        return {"ok": False, "err": str(e)}


# ----------- 配置文件处理 -----------
CONFIG_FILE = Path(__file__).with_name("config.json")

def load_config():
    """加载配置文件，如果不存在则返回空配置"""
    if CONFIG_FILE.exists():
        try:
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except:
            return {}
    return {}

def save_config(config: dict):
    """保存配置到文件"""
    with open(CONFIG_FILE, "w", encoding="utf-8") as f:
        json.dump(config, f, ensure_ascii=False, indent=2)


# ----------- SSH 工具函数 -----------
def create_ssh_client(host: str, port: int, user: str, pwd: str, timeout: int = 10) -> paramiko.SSHClient:
    """创建并连接SSH客户端"""
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=host, port=port, username=user, password=pwd, timeout=timeout)
    return ssh

async def execute_command_and_stream_output(websocket: WebSocket, chan, cmd: str, timeout: int = 10):
    """执行命令并通过WebSocket流式传输输出"""
    try:
        await websocket.send_json({"type": "msg", "data": f"执行：{cmd}"})
        # 在函数中执行命令
        chan.exec_command(cmd)
        
        # 获取实时输出
        # 设置超时，避免无限等待
        chan.settimeout(timeout)
        
        while not chan.exit_status_ready():
            if chan.recv_ready():
                chunk = chan.recv(4096).decode(errors="ignore")
                if not chunk:
                    break
                # 修改: 使用JSON格式发送消息
                await websocket.send_json({"type": "msg", "data": chunk})
            if chan.recv_stderr_ready():
                error_chunk = chan.recv_stderr(4096).decode(errors="ignore")
                if error_chunk:
                    # 修改: 使用JSON格式发送错误消息
                    await websocket.send_json({"type": "msg", "data": error_chunk})
            await asyncio.sleep(0.05)
        
        # 处理剩余输出
        while chan.recv_ready():
            chunk = chan.recv(4096).decode(errors="ignore")
            if not chunk:
                break
            # 修改: 使用JSON格式发送消息
            await websocket.send_json({"type": "msg", "data": chunk})
            
        exit_code = chan.recv_exit_status()
        return exit_code
    except Exception as e:
        # 修改: 使用JSON格式发送错误
        await websocket.send_json({"type": "msg", "data": f"[ERROR] {e}"})
        raise


# ----------- API -----------
@app.post("/api/test")
def api_test_ssh(payload: Dict[str, Any]):
    # 将前端传递的 ip 参数转换为 host 参数
    if 'ip' in payload:
        payload['host'] = payload.pop('ip')
    
    # 逐个提取参数而不是使用**解包
    host = payload.get('host')
    port = payload.get('port')
    user = payload.get('user')
    pwd = payload.get('pwd')
    
    # 检查必需参数
    if not all([host, port, user, pwd]):
        return {"ok": False, "err": "缺少必需参数"}
    
    return test_ssh(host, port, user, pwd)


@app.post("/api/images")
def api_images(payload: Dict[str, Any]):
    # 将前端传递的 ip 参数转换为 host 参数
    if 'ip' in payload:
        payload['host'] = payload.pop('ip')
        
    # 逐个提取参数而不是使用**解包
    host = payload.get('host')
    port = payload.get('port')
    user = payload.get('user')
    pwd = payload.get('pwd')
    
    # 检查必需参数
    if not all([host, port, user, pwd]):
        return {"ok": False, "err": "缺少必需参数"}
    
    return list_docker_images(host, port, user, pwd)

# 添加配置读取和保存的API
@app.get("/api/config")
def api_get_config():
    return load_config()

@app.post("/api/config")
def api_save_config(payload: Dict[str, Any]):
    save_config(payload)
    return {"ok": True}


@app.post("/api/checkGpuStatus")
def check_gpu_status():
    # TODO: 实现GPU状态检查逻辑
    return {"ok": True}


# ----------- WebSocket 实时日志 -----------
# 放在 main.py 里，覆盖原来的 @app.websocket("/ws/run")
@app.websocket("/ws/run", name="run")
async def websocket_run(websocket: WebSocket):
    await websocket.accept()

    # 1. 接收前端发来的配置与命令
    msg = await websocket.receive_text()
    data = json.loads(msg)          # {"host":..,"port":..,"user":..,"pwd":..,"cmd":"docker run ..."}
    host, port, user, pwd, cmd = (
        data["host"],
        data["port"],
        data["user"],
        data["pwd"],
        data["cmd"],
    )

    ssh = None
    try:
        # 2. Paramiko 连接并执行
        ssh = create_ssh_client(host, port, user, pwd, timeout=10)
        chan = ssh.get_transport().open_session()
        exit_code = await execute_command_and_stream_output(websocket, chan, cmd, 10)
        await websocket.send_json({"type": "exit", "data": exit_code})
    except Exception as e:
        # 修改: 使用JSON格式发送错误
        await websocket.send_json({"type": "msg", "data": f"[ERROR] {e}"})
        await websocket.send_json({"type": "exit", "data": -1})
    finally:
        if ssh:
            try:
                ssh.close()
            except:
                pass
        await websocket.close()


# 新增WebSocket端点用于编译操作
@app.websocket("/ws/build", name="build")
async def websocket_build(websocket: WebSocket):
    await websocket.accept()
    
    ssh = None
    try:
        # 从配置文件读取编译服务器配置，而不是从前端接收
        config = load_config()
        build_config = config.get("build", {})
        host = build_config.get("host") or build_config.get("ip")
        port = build_config.get("port", 22)
        user = build_config.get("user")
        pwd = build_config.get("pwd")
        code_source = build_config.get("codeSource", "remote")
        code_path = build_config.get("codePath", "")
        git_repo = build_config.get("gitRepo", "")
        git_branch = build_config.get("gitBranch", "main")
        
        if not all([host, user, pwd]):
            # 修改: 使用JSON格式发送错误消息
            await websocket.send_json({"type": "msg", "data": "[编译] 错误: 编译服务器配置不完整\n"})
            await websocket.close()
            return
            
        # 根据代码来源执行不同逻辑
        if code_source == "remote":
            if not code_path:
                # 修改: 使用JSON格式发送错误消息
                await websocket.send_json({"type": "msg", "data": "[编译] 错误: 代码路径未配置\n"})
                await websocket.close()
                return
            cmd = f"cd {code_path} && ls -l"
            # 修改: 使用JSON格式发送消息
            await websocket.send_json({"type": "msg", "data": f"[编译] 在远程路径 {code_path} 执行命令: ls -l\n"})
        else:
            # 本地代码上传逻辑（这里简化处理，实际可能需要实现文件上传）
            # 检查Git仓库和分支配置
            if not git_repo:
                # 修改: 使用JSON格式发送错误消息
                await websocket.send_json({"type": "msg", "data": "[编译] 错误: Git仓库地址未配置\n"})
                await websocket.close()
                return
                
            cmd = f"git clone {git_repo} temp_repo && cd temp_repo && git checkout {git_branch} && ls -l"
            # 修改: 使用JSON格式发送消息
            await websocket.send_json({"type": "msg", "data": f"[编译] 克隆Git仓库 {git_repo}，分支: {git_branch}\n"})
        
        # SSH连接并执行编译命令
        ssh = create_ssh_client(host, port, user, pwd, timeout=10)
        chan = ssh.get_transport().open_session()
        exit_code = await execute_command_and_stream_output(websocket, chan, cmd, 10)
        await websocket.send_json({"type": "exit", "data": exit_code})
        
    except Exception as e:
        # 修改: 使用JSON格式发送错误
        await websocket.send_json({"type": "msg", "data": f"[编译] 错误: {e}\n"})
        await websocket.send_json({"type": "exit", "data": -1})
    finally:
        if ssh:
            try:
                ssh.close()
            except:
                pass
        await websocket.close()


# 新增WebSocket端点用于运行操作
@app.websocket("/ws/run-command", name="run_command")
async def websocket_run_command(websocket: WebSocket):
    await websocket.accept()
    
    ssh = None
    try:
        # 接收前端发送的命令
        data = await websocket.receive_json()
        command = data.get("command", "ls -lh")
        
        # 从配置文件读取运行服务器配置，而不是从前端接收
        config = load_config()
        run_config = config.get("run", {})
        host = run_config.get("host") or run_config.get("ip")
        port = run_config.get("port", 22)
        user = run_config.get("user")
        pwd = run_config.get("pwd")
        
        if not all([host, user, pwd]):
            # 修改: 使用JSON格式发送错误消息
            await websocket.send_json({"type": "msg", "data": "[运行] 错误: 运行服务器配置不完整\n"})
            await websocket.close()
            return
        
        # 修改: 使用JSON格式发送消息
        await websocket.send_json({"type": "msg", "data": f"[运行] 开始执行 {command} 命令...\n"})
        
        # SSH连接并执行运行命令
        ssh = create_ssh_client(host, port, user, pwd, timeout=10)
        chan = ssh.get_transport().open_session()
        exit_code = await execute_command_and_stream_output(websocket, chan, command, 10)
        await websocket.send_json({"type": "exit", "data": exit_code})
        
    except Exception as e:
        # 修改: 使用JSON格式发送错误
        await websocket.send_json({"type": "msg", "data": f"[运行] 错误: {e}\n"})
        await websocket.send_json({"type": "exit", "data": -1})
    finally:
        if ssh:
            try:
                ssh.close()
            except:
                pass
        await websocket.close()


# 新增WebSocket端点用于编译命令执行
@app.websocket("/ws/compile-command", name="compile_command")
async def websocket_compile_command(websocket: WebSocket):
    await websocket.accept()
    
    ssh = None
    try:
        # 接收前端发送的命令
        data = await websocket.receive_json()
        command = data.get("command", "ls -l")
        
        # 从配置文件读取编译服务器配置，而不是从前端接收
        config = load_config()
        build_config = config.get("build", {})
        host = build_config.get("host") or build_config.get("ip")
        port = build_config.get("port", 22)
        user = build_config.get("user")
        pwd = build_config.get("pwd")
        code_source = build_config.get("codeSource", "remote")
        code_path = build_config.get("codePath", "")
        git_repo = build_config.get("gitRepo", "")
        git_branch = build_config.get("gitBranch", "main")
        
        if not all([host, user, pwd]):
            # 修改: 使用JSON格式发送错误消息
            await websocket.send_json({"type": "msg", "data": "[编译] 错误: 编译服务器配置不完整\n"})
            await websocket.close()
            return
            
        # 根据代码来源执行不同逻辑
        if code_source == "remote":
            if not code_path:
                # 修改: 使用JSON格式发送错误消息
                await websocket.send_json({"type": "msg", "data": "[编译] 错误: 代码路径未配置\n"})
                await websocket.close()
                return
            full_command = f"cd {code_path} && {command}"
            # 修改: 使用JSON格式发送消息
            await websocket.send_json({"type": "msg", "data": f"[编译] 在远程路径 {code_path} 执行命令: {command}\n"})
        else:
            # 本地代码上传逻辑（这里简化处理，实际可能需要实现文件上传）
            # 检查Git仓库和分支配置
            if not git_repo:
                # 修改: 使用JSON格式发送错误消息
                await websocket.send_json({"type": "msg", "data": "[编译] 错误: Git仓库地址未配置\n"})
                await websocket.close()
                return
                
            full_command = f"git clone {git_repo} temp_repo && cd temp_repo && git checkout {git_branch} && {command}"
            # 修改: 使用JSON格式发送消息
            await websocket.send_json({"type": "msg", "data": f"[编译] 克隆Git仓库 {git_repo}，分支: {git_branch}，执行命令: {command}\n"})
        
        # SSH连接并执行编译命令
        ssh = create_ssh_client(host, port, user, pwd, timeout=10)
        chan = ssh.get_transport().open_session()
        exit_code = await execute_command_and_stream_output(websocket, chan, full_command, 10)
        await websocket.send_json({"type": "exit", "data": exit_code})
        
    except Exception as e:
        # 修改: 使用JSON格式发送错误
        await websocket.send_json({"type": "msg", "data": f"[编译] 错误: {e}\n"})
        await websocket.send_json({"type": "exit", "data": -1})
    finally:
        if ssh:
            try:
                ssh.close()
            except:
                pass
        await websocket.close()


# 新增WebSocket端点用于本地命令执行
@app.websocket("/ws/local-command", name="local_command")
async def websocket_local_command(websocket: WebSocket):
    await websocket.accept()
    
    try:
        # 接收前端发送的命令
        data = await websocket.receive_json()
        command = data.get("command", "echo 'Hello World'")
        
        await websocket.send_json({"type": "msg", "data": f"[本地] 开始执行命令: {command}\n"})
        
        # 创建子进程执行本地命令
        process = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        # 读取输出
        while True:
            try:
                # 等待输出或超时
                stdout_line = await asyncio.wait_for(process.stdout.readline(), timeout=1.0)
                if stdout_line:
                    await websocket.send_json({"type": "msg", "data": stdout_line.decode()})
                else:
                    break
            except asyncio.TimeoutError:
                # 检查进程是否仍在运行
                if process.returncode is not None:
                    break
                continue
        
        # 等待进程完成
        stdout, stderr = await process.communicate()
        
        # 发送剩余输出
        if stdout:
            await websocket.send_json({"type": "msg", "data": stdout.decode()})
        if stderr:
            await websocket.send_json({"type": "msg", "data": f"错误: {stderr.decode()}"})
        
        # 发送退出码
        await websocket.send_json({"type": "exit", "data": process.returncode})
        
    except Exception as e:
        await websocket.send_json({"type": "msg", "data": f"[本地] 错误: {e}\n"})
        await websocket.send_json({"type": "exit", "data": -1})
    finally:
        await websocket.close()

# 辅助函数：执行命令并流式输出结果

# 把 /assets 目录挂到根路径，并把 index.html 设为默认首页
# 注意：需要在所有API路由注册完成后再挂载静态文件，否则会拦截API请求
app.mount("/", StaticFiles(directory=Path(__file__).with_name("assets"), html=True), name="static")

# 添加启动入口点
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)